import os
import sys
import yaml
import sherpa_onnx
import sounddevice as sd
import websockets
import queue
import asyncio
import threading
# 用一个线程安全的队列来传递图像
message_queue = queue.Queue()

def create_recognizer():
    recognizer = sherpa_onnx.OnlineRecognizer.from_transducer(
        tokens="./models/sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20/tokens.txt",
        encoder="./models/sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20/encoder-epoch-99-avg-1.onnx",
        decoder= "./models/sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20/decoder-epoch-99-avg-1.onnx",
        joiner="./models/sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20/joiner-epoch-99-avg-1.onnx",
        num_threads=1,
        sample_rate=16000,
        feature_dim=80,
        enable_endpoint_detection=True,
        rule1_min_trailing_silence=0.4,
        rule2_min_trailing_silence=0.3,
        rule3_min_utterance_length=300,  # it essentially disables this rule
        decoding_method= "greedy_search",
        provider ="cpu",
        hotwords_score=1.5,
        blank_penalty=0,
    )
    return recognizer


def process_asr():
    devices = sd.query_devices()
    if len(devices) == 0:
        print("No microphone devices found")
        sys.exit(0)
    # default_input_device_idx = sd.default.device[0]
    # print(f'Use default device: {devices[default_input_device_idx]["name"]}')

    recognizer = create_recognizer()
  
    sample_rate = 48000
    samples_per_read = int(0.1 * sample_rate)  # 0.1 second = 100 ms
    stream = recognizer.create_stream()
    # display = sherpa_onnx.Display()

    with sd.InputStream(channels=1, dtype="float32", samplerate=sample_rate) as s:
        while True:
            samples, _ = s.read(samples_per_read)  # a blocking read
            samples = samples.reshape(-1)
            stream.accept_waveform(sample_rate, samples)
            while recognizer.is_ready(stream):
                recognizer.decode_stream(stream)
            is_endpoint = recognizer.is_endpoint(stream)
            result = recognizer.get_result(stream)
            # display.update_text(result)
            # display.display()
            if is_endpoint:
                if result:
                    message_queue.put(result)
                    # display.finalize_current_sentence()
                    # display.display()
                recognizer.reset(stream)

async def echo(websocket):
    while not message_queue.empty():
        message =message_queue.get()  # 获取到后立刻清空队列
    while True:
        if not message_queue.empty():
            message = message_queue.get()
            await websocket.send(message)
        await asyncio.sleep(0.02)  # 避免占用过多 CPU

async def main():
    async with websockets.serve(echo, "0.0.0.0", 5000):
        await asyncio.Future()  # run forever

# 将处理逻辑封装到线程中
process_thread = threading.Thread(target=process_asr)
process_thread.start()

asyncio.run(main())